use std::str;
use std::sync::{Arc, Mutex};

use anyhow::{bail, ensure, Result};
use async_h1::server::decode;
use async_std::future::{timeout, TimeoutError};
use async_std::io::{self, BufReader, ReadExt};
use async_std::net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener, TcpStream};
use async_std::prelude::*;
use async_std::task;
use async_tls::TlsAcceptor;

use async_h1::client::Encoder;
use std::time::Duration;
use tls::{load_config, TLStream};

const END_0: u8 = b'\0';
const LAMBDA_SERVER: &str = "127.0.0.1:8081";
const TIMEOUT: Duration = Duration::from_secs(3);

fn main() {
    pretty_env_logger::init_timed();
    if let Err(e) = task::block_on(run()) {
        log::error!("run fail: {}", e)
    }
}

async fn run() -> Result<()> {
    let listener = TcpListener::bind("127.0.0.1:8080").await?;

    // We create one TLSAcceptor around a shared configuration.
    // Cloning the acceptor will not clone the configuration.
    let config = load_config()?;
    let acceptor = TlsAcceptor::from(Arc::new(config));

    loop {
        let (stream, _) = listener.accept().await?;
        let acceptor = acceptor.clone();
        task::spawn(async move {
            if let Err(err) = handle(&acceptor, stream).await {
                log::error!("handle fail: [{:?}]", err)
            }
        });
    }
}

async fn handle(acceptor: &TlsAcceptor, proxy: TcpStream) -> Result<()> {
    let mut reader = BufReader::new(proxy.clone());
    let mut buf = [0_u8; 8];
    reader.read_exact(&mut buf).await?;
    log::info!("head: {:?}", buf);
    ensure!(buf.len() >= 8, "长度不足");
    ensure!(buf[0] == v4::VERSION, "不支持的Version: [{}]", buf[0]);
    ensure!(buf[1] == v4::CMD_CONNECT, "不支持的Command: [{}]", buf[1]);

    // Read until \0
    let mut msg = Vec::new();
    reader.read_until(END_0, &mut msg).await?;

    // Read Ipv4 Addr
    let port = ((buf[2] as u16) << 8) | (buf[3] as u16);
    let ip = Ipv4Addr::new(buf[4], buf[5], buf[6], buf[7]);
    let ip_port = format!("{}.{}.{}.{}:{}", buf[4], buf[5], buf[6], buf[7], port);

    // Connect
    let octets = ip.octets();
    let remote = if octets[0] == 0 && octets[1] == 0 && octets[2] == 0 && octets[3] != 0 {
        log::debug!("socks4a: {}:{}", ip, port);

        // SOCKS4a
        reader.read_until(END_0, &mut msg).await?;
        let host_name = str::from_utf8(&*msg)?;
        log::debug!("host_name: {}", host_name);

        TcpStream::connect(host_name).await?
    } else {
        log::debug!("socks4: {}:{}", ip, port);

        // SOCKS4
        TcpStream::connect(SocketAddr::new(IpAddr::V4(ip), port)).await?
    };

    // Generate Response
    let mut resp = bytes::BytesMut::new();
    let addr = remote.local_addr()?;
    // let addr = proxy.local_addr()?;
    resp.extend_from_slice(
        &[
            0_u8,
            v4::CMD_ALLOW,
            (addr.port() >> 8) as u8,
            addr.port() as u8,
        ][..],
    );

    if let IpAddr::V4(ip) = addr.ip() {
        resp.extend(ip.octets().iter());
    } else {
        bail!("addr is not IpAddr::V4, [{:?}]", addr)
    }

    // Write Response
    proxy.clone().write_all(&*resp).await?;

    // https Handshake and read http request from browser
    let tls_proxy = acceptor.accept(proxy).await?;
    let mut tls_proxy = TLStream::new(Arc::new(Mutex::new(tls_proxy)));

    let (mut req, mut body) = match timeout(TIMEOUT, decode(tls_proxy.clone())).await {
        Ok(Ok(Some(r))) => r,
        Ok(Ok(None)) | Err(TimeoutError { .. }) => bail!("Timeout!"), /* EOF or timeout */
        Ok(Err(e)) => anyhow::bail!("{}", e),
    };
    io::copy(&mut body, &mut io::sink()).await?;

    // async_h1::accept(tls_proxy, |mut req| async {
    //     // req.ext_mut().insert(SocketAddr::new(IpAddr::V4(ip), port));
    //     // test_tls(req.clone(), ip, port).await?;
    //
    //     // Connect lambda and send http request
    //     let remote = TcpStream::connect(LAMBDA_SERVER).await?;
    //
    //     //TODO 加密req
    //     log::info!("req: {:?}", req);
    //
    //     //调用lambda
    //     let res = async_h1::connect(remote, req).await?;
    //     Ok(res)
    // })
    // .await
    // .map_err(|e| anyhow::anyhow!("{}", e))?;

    // Connect lambda and send http request
    let mut remote = TcpStream::connect(LAMBDA_SERVER).await?;

    //TODO 加密req
    log::info!("req: {:?}", req);
    req.append_header("ip_port", ip_port);

    let mut req = Encoder::new(req);
    io::copy(&mut req, &mut remote).await?;
    log::debug!("send over");
    io::copy(&mut remote, &mut tls_proxy).await?;
    log::debug!("receive over");

    Ok(())
}

mod v4 {
    pub const VERSION: u8 = 4;
    pub const CMD_CONNECT: u8 = 1;
    pub const CMD_ALLOW: u8 = 90_u8;
}

// async fn test_tls(req: Request, ip: Ipv4Addr, port: u16) -> Result<()> {
//     // let req = Request::new(Method::Get, req.url());
//     log::info!("begin test: {:?}", req);
//     // let remote = TcpStream::connect(req.host().unwrap()).await?;
//     let remote = TcpStream::connect(SocketAddr::new(IpAddr::V4(ip), port)).await?;
//
//     log::info!("connect {}", req.host().unwrap());
//     let tls_remote = TlsConnector::default()
//         .connect(req.host().unwrap(), remote)
//         .await?;
//     log::info!("send request");
//     let res = async_h1::connect(tls_remote, req)
//         .await
//         .map_err(|e| anyhow::anyhow!("connect fail {}", e))?;
//     log::info!("receive response: {:#?}", res);
//
//     Ok(())
// }
